from rknn.api import RKNN

# 标准版本：使用640x640正方形输入（推荐）
# 这与你的YOLO训练时的imgsz=640设置一致
rknn = RKNN(verbose=True)

# 配置RKNN转换参数
rknn.config(
    mean_values=[[0, 0, 0]],           # 均值归一化（因为在预处理中已经做了/255）
    std_values=[[255, 255, 255]],     # 标准差归一化
    # quantized_dtype='w8a8',          # 如果要量化，取消注释
    # quantized_algorithm='normal',    # 如果要量化，取消注释
    optimization_level=3,              # 最高优化级别
    target_platform='rk3588',         # 目标平台
    output_optimize=True,              # 输出优化
)

# 加载ONNX模型 - 使用标准640x640尺寸
print("加载ONNX模型...")
ret = rknn.load_onnx(
    model='best.onnx',
    inputs=['images'],
    input_size_list=[[1, 3, 640, 640]]  # 标准YOLOv8输入尺寸（正方形），与训练时一致
)

if ret != 0:
    print('加载模型失败！请检查ONNX文件路径和格式')
    exit(ret)

print("构建RKNN模型...")
ret = rknn.build(do_quantization=False)  # 不进行量化，保持FP16精度

if ret != 0:
    print('构建模型失败！')
    exit(ret)

print("导出RKNN模型...")
ret = rknn.export_rknn('wuzi_640x640.rknn')

if ret != 0:
    print('导出RKNN失败！')
    exit(ret)

print('✅ RKNN模型转换成功！')
print('输出文件: wuzi_640x640.rknn')
print('输入格式: [1, 3, 640, 640] (NCHW)')
print('建议：')
print('1. 使用wuzi_640x640.rknn替换原来的wuzi.rknn')
print('2. 推理代码将自动使用letterbox预处理保持宽高比')
print('3. 这种方式与你的训练设置完全一致，推荐使用')

rknn.release() 